/*
 * LoopCountInsertion.cpp
 *
 *  Created on: Oct 30, 2011
 *      Author: kula85
 */

#define DEBUG_TYPE "delaythread"
#include "llvm/Pass.h"
#include "llvm/Function.h"
#include "llvm/BasicBlock.h"
#include "llvm/Instructions.h"
#include "llvm/DerivedTypes.h"
#include "llvm/Module.h"
#include "llvm/GlobalValue.h"
#include "llvm/Constants.h"
#include "llvm/GlobalVariable.h"
#include "llvm/ADT/VectorExtras.h"
#include "llvm/ADT/ArrayRef.h"
#include "llvm/ADT/StringRef.h"
#include "llvm/ADT/SmallPtrSet.h"
#include <cstdio>
#include <vector>
#include <string>
#include "llvm/Transforms/Scalar.h"
#include "llvm/Analysis/DebugInfo.h"
#include "llvm/Analysis/LoopInfo.h"
#include "llvm/Support/InstIterator.h"
#include "llvm/Support/raw_ostream.h"
#include "llvm/Support/IRBuilder.h"
#include "llvm/Metadata.h"
using namespace llvm;


namespace {
struct LoopCountInsertion : public ModulePass {
  static char ID;
  LoopCountInsertion() : ModulePass(ID) {
    initializeLoopCountInsertionPass(*PassRegistry::getPassRegistry());
  }

  virtual bool runOnModule(Module &M) {
    unsigned LoopCounterIndex = 0;
    Constant *posvalue;
    bool hasmain = false;
    bool haspapikernel = false;

    // if source file contains no function. return
    if(M.getFunctionList().size() == 0) return false;

    // test for hasmain & haspapikernel
    for(Module::iterator MI=M.begin(), ME=M.end(); MI != ME; MI++) {
      if (MI->getName() == "main") {
        hasmain = true;
      }
      if (MI->getName() == "PapiKernel") {
        haspapikernel = true;
      }
    }

    if(hasmain && !haspapikernel) return false;

    // run this pass only once
    GlobalVariable *GV = M.getNamedGlobal(StringRef("__dt_pos"));
    //if(GV && GV->getLinkage() == GlobalValue::ExternalLinkage) return false;
    //if(GV && !haspapikernel) return false;
    if(GV && !haspapikernel) return false;
    if(GV && haspapikernel && M.getFunction("strcpy")) return false;
    // if file is our main, return


    IRBuilder<> Builder(&*M.begin()->begin());

    //add current filename to global
    NamedMDNode *CU_Nodes = M.getNamedMetadata("llvm.dbg.cu");
    assert( CU_Nodes && "No NamedMDNode for Module. Forgot to use '-g' flag?");
    assert( CU_Nodes->getNumOperands() == 1 &&
        "Moule has more than 1 CompileUnit metadata?");
    DICompileUnit CU(CU_Nodes->getOperand(0));
    StringRef Dir = CU.getDirectory();
    StringRef File = CU.getFilename();
    Value *dirnamestr = Builder.CreateGlobalString(Dir, Twine("__dt_dirname"));
    Value *filenamestr = Builder.CreateGlobalString(File, Twine("__dt_filename"));

    // struct type already exists
    StructType *postype = M.getTypeByName("struct.__dt_position");
    // add struct type to store position info
    if(!postype) {
      std::vector<Type*> infotypes = make_vector<Type*>(
          ArrayType::get(IntegerType::getInt8Ty(M.getContext()),256), //dirname
          ArrayType::get(IntegerType::getInt8Ty(M.getContext()),256), //filename
          IntegerType::getInt32Ty(M.getContext()), //loop line number
          IntegerType::getInt32Ty(M.getContext()),0); // iteration

      postype = StructType::create(M.getContext(),
          makeArrayRef(infotypes), StringRef("struct.__dt_position"));
    }

    // add pos var
    if(hasmain) {
      if(haspapikernel) {
        posvalue = M.getNamedGlobal("__dt_pos");
      } else{
        // add struct init
        ConstantInt *cinit = ConstantInt::get(Type::getInt32Ty(M.getContext()), 0);
        Constant *sinit = ConstantArray::get(M.getContext(), "");

        std::vector<Constant*> svc;
        svc.push_back(sinit);
        svc.push_back(sinit);
        svc.push_back(cinit);
        svc.push_back(cinit);

        Constant *init = ConstantStruct::get(postype, makeArrayRef(svc));

        posvalue = new GlobalVariable(M, postype,
            false, GlobalValue::CommonLinkage, init, "__dt_pos");
      }
    } else
      posvalue = M.getOrInsertGlobal(StringRef("__dt_pos"), postype);

    posvalue.setAlignment(16);

    // add declaration for strcpy
    Function *strcpy_func = M.getFunction("strcpy");
    if(!strcpy_func) {
      std::vector<Type*> strcpy_2args = make_vector<Type*>(
          Type::getInt8PtrTy(M.getContext()),
          Type::getInt8PtrTy(M.getContext()),0);

      FunctionType *strcpy_funcTy = FunctionType::get(
          /*Result=*/Type::getInt8PtrTy(M.getContext()),
          /*Params=*/strcpy_2args,
          /*isVarArg=*/false);

      strcpy_func = Function::Create(
          /*Type=*/strcpy_funcTy,
          /*Linkage=*/GlobalValue::ExternalLinkage,
          /*Name=*/"strcpy", &M); // (external, no body)

      //strcpy_func->addAttribute(2, Attribute::NoCapture);
      strcpy_func->addFnAttr(Attribute::NoUnwind);
    }

    // add declaration for pthread_self
    Function *pself_func = M.getFunction("pthread_self");
    if(!pself_func) {
      FunctionType *pself_funcTy = FunctionType::get(
          /*Result=*/Type::getInt64PtrTy(M.getContext()),
          /*isVarArg=*/false);

      pself_func = Function::Create(
          /*Type=*/pself_funcTy,
          /*Linkage=*/GlobalValue::ExternalLinkage,
          /*Name=*/"pthread_self", &M); // (external, no body)

      pself_func->addFnAttr(Attribute::NoUnwind);
      pself_func->addFnAttr(Attribute::ReadNone);
    }

    // get all loops in this module
    SmallPtrSet<Loop*, 5> AllLoops;
    //SmallPtrSet<LoopInfo*, 5> AllLoopInfos;
    for(Module::iterator MI=M.begin(), ME=M.end(); MI != ME; MI++) {
      //only a declaration?
      if(MI->isDeclaration()) continue;
      //function with no name?
      if(MI->getName().size() == 0) continue;

      Value *Idxs1[] = {
          ConstantInt::get(Type::getInt64Ty(M.getContext()), 0),
          ConstantInt::get(Type::getInt32Ty(M.getContext()), 0),
          ConstantInt::get(Type::getInt64Ty(M.getContext()), 0)
      };

      Value *Idxs2[] = {
          ConstantInt::get(Type::getInt64Ty(M.getContext()), 0),
          ConstantInt::get(Type::getInt32Ty(M.getContext()), 1),
          ConstantInt::get(Type::getInt64Ty(M.getContext()), 0)
      };

      Value *Idxs3[] = {
          ConstantInt::get(Type::getInt64Ty(M.getContext()), 0),
          ConstantInt::get(Type::getInt64Ty(M.getContext()), 0)
      };

      Value *Idxs4[] = {
          ConstantInt::get(Type::getInt64Ty(M.getContext()), 0),
          ConstantInt::get(Type::getInt32Ty(M.getContext()), 2)
      };

      Value *Idxs5[] = {
          ConstantInt::get(Type::getInt64Ty(M.getContext()), 0),
          ConstantInt::get(Type::getInt32Ty(M.getContext()), 3)
      };


      //monitoring thread don't need to run this pass, but need to change
      //fprintf to output position info
      if(MI->getName().equals(StringRef("PapiKernel"))) {
//        CallInst* temp = 0;
//        Value *fprtstr = Builder.CreateGlobalString(
//                  StringRef("%lld\t%lld\t%lld\t%s\t%s\t%u\t%u\n"));
//        assert(isa<GlobalVariable>(fprtstr) &&
//                      "fprtstr is not global variable?");
//        GlobalVariable *gv = cast<GlobalVariable>(fprtstr);
//        gv->setLinkage(GlobalVariable::PrivateLinkage);
//        gv->setAlignment(1);
//        std::vector<Value*> params;
//
//        for(inst_iterator ii = inst_begin(*MI),
//                        ie = inst_end(*MI); ii != ie; ++ii) {
//          if (CallInst* callInst = dyn_cast<CallInst>(&*ii)) {
//            Function* func = callInst->getCalledFunction();
//            if ((func->getName()).startswith(StringRef("fprintf"))) {
//              for(unsigned i = 0; i < callInst->getNumArgOperands(); i++) {
//                  if(i==1) params.push_back(Builder.
//                      CreateInBoundsGEP(gv, makeArrayRef(Idxs3)));
//                  else params.push_back(callInst->getArgOperand(i));
//              }
//
//              Builder.SetInsertPoint(callInst);
//
//              LoadInst *loadpos0 = Builder.CreateLoad(
//                  Builder.CreateInBoundsGEP(posvalue, makeArrayRef(Idxs1)));
//              LoadInst *loadpos1 = Builder.CreateLoad(
//                  Builder.CreateInBoundsGEP(posvalue, makeArrayRef(Idxs2)));
//              LoadInst *loadpos2 = Builder.CreateLoad(
//                  Builder.CreateInBoundsGEP(posvalue, makeArrayRef(Idxs4)));
//              LoadInst *loadpos3 = Builder.CreateLoad(
//                  Builder.CreateInBoundsGEP(posvalue, makeArrayRef(Idxs5)));
//
//              params.push_back(loadpos0);
//              params.push_back(loadpos1);
//              params.push_back(loadpos2);
//              params.push_back(loadpos3);
//
//              //Builder.SetInsertPoint(callInst);
//              Builder.CreateCall(func, makeArrayRef(params));
//              temp = callInst;
//            }
//          }
//        }
//
//        temp->eraseFromParent();
//
        continue;
      }

      LoopInfo &LI = getAnalysis<LoopInfo>(*MI);
      for(Function::iterator FI=MI->begin(), FE=MI->end(); FI != FE; FI++) {
        // if L is null or L is already processed, goto next BB
        Loop *L = LI.getLoopFor(&*FI);
        if(L == 0) continue;
        if(!AllLoops.insert(L)) continue;

        // add struct type for pos
        unsigned Line = 0;
        GlobalVariable *loopcounter = new GlobalVariable(M,
            IntegerType::getInt32Ty(Builder.getContext()),
            false, GlobalValue::InternalLinkage, Builder.getInt32(0),
            Twine("__dt_loopcounter")+Twine(LoopCounterIndex));
        loopcounter->setAlignment(4);

        // add loop counter in loop body
        BasicBlock *loopbody = L->getHeader()->getTerminator()->getSuccessor(0);

        Builder.SetInsertPoint(loopbody->getFirstNonPHI());
        LoadInst *loadindex = Builder.CreateLoad(loopcounter);
        //loadindex->setVolatile(1);
        Value *addinst = Builder.CreateNSWAdd(loadindex, Builder.getInt32(1));
        StoreInst *storeinst = Builder.CreateStore(addinst, loopcounter);
        //storeinst->setVolatile(1);

        // add two strcpy calls for dirname&filename
        Builder.CreateCall2(strcpy_func,
            Builder.CreateInBoundsGEP(posvalue, makeArrayRef(Idxs1)),
            Builder.CreateInBoundsGEP(dirnamestr, makeArrayRef(Idxs3)));

        Builder.CreateCall2(strcpy_func,
            Builder.CreateInBoundsGEP(posvalue, makeArrayRef(Idxs2)),
            Builder.CreateInBoundsGEP(filenamestr, makeArrayRef(Idxs3)));

        MDNode *N = L->getHeader()->getTerminator()
                                           ->getPrevNode()->getMetadata("dbg");
        //assert(N && "No Line Info for loophead. Forget to use '-g' flag?\n");
        DILocation Loc(N);
        Line = Loc.getLineNumber();

        storeinst = Builder.CreateStore(ConstantInt::get(
            Type::getInt32Ty(M.getContext()), Line),
            Builder.CreateInBoundsGEP(posvalue, makeArrayRef(Idxs4)));
        storeinst->setVolatile(1);

        loadindex = Builder.CreateLoad(loopcounter);
        Builder.CreateStore(loadindex,
            Builder.CreateInBoundsGEP(posvalue, makeArrayRef(Idxs5)));
        loadindex->setVolatile(1);
      }
    }

    return true;
  }

  // We don't modify the program, so we preserve all analyses
  void getAnalysisUsage(AnalysisUsage &AU) const {
    AU.setPreservesCFG();
    AU.addRequired<LoopInfo>();
  }
};
}

char LoopCountInsertion::ID = 0;
INITIALIZE_PASS_BEGIN(LoopCountInsertion, "loop-count-insertion",
    "Pass to add loop profiling counter to know "
    "which iteration we are in. "
    "This pass also generate a file "
    "which contains info about which counter is for "
    "which loop(line no is for cmp inst).", false, false)
INITIALIZE_PASS_DEPENDENCY(LoopInfo)
INITIALIZE_PASS_END(LoopCountInsertion, "loop-count-insertion",
    "Pass to add loop profiling counter to know "
    "which iteration we are in. "
    "This pass also generate a file "
    "which contains info about which counter is for "
    "which loop(line no is for cmp inst).", false, false)

    Pass *llvm::createLoopCountInsertionPass() {
  return new LoopCountInsertion();
}
